package com.example.weather.utils.Tree;

import org.apache.poi.ss.formula.functions.T;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树形结构工具类
 */
public class TreeUtil {

    /**
     * 构建树形结构
     * @param list 原始数据列表
     * @param rootParentId 根节点的父ID值
     * @return 树形结构列表
     */
    public static <T, N extends TreeNode<T, N>> List<N> buildTree(List<N> list, T rootParentId) {
        // 参数校验
        if (list == null || list.isEmpty()) {
            return new ArrayList<>();
        }

        // 按ID分组
        Map<T, N> nodeMap = list.stream()
                .collect(Collectors.toMap(TreeNode::getId, Function.identity()));

        List<N> roots = new ArrayList<>();

        for (N node : list) {
            T parentId = node.getParentId();
            if (Objects.equals(parentId, rootParentId)) {
                roots.add(node);
            } else {
                N parent = nodeMap.get(parentId);
                if (parent != null) {
                    if (parent.getChildren() == null) {
                        parent.setChildren(new ArrayList<>());
                    }
                    parent.getChildren().add(node);
                }
            }
        }

        return roots;
    }

    /**
     * 构建树形结构（自动识别根节点）
     * @param list 原始数据列表
     * @return 树形结构列表
     */
    public static <T, N extends TreeNode<T, N>> List<N> buildTree(List<N> list) {
        if (list == null || list.isEmpty()) {
            return new ArrayList<>();
        }

        // 找出所有可能的父节点ID
        List<T> allIds = list.stream().map(TreeNode::getId).collect(Collectors.toList());
        
        // 找出根节点（父节点不在所有ID中的节点）
        List<N> roots = list.stream()
                .filter(node -> !allIds.contains(node.getParentId()))
                .collect(Collectors.toList());

        // 如果没找到根节点，尝试找出parentId为null或特定值的节点
        if (roots.isEmpty()) {
            roots = list.stream()
                    .filter(node -> node.getParentId() == null)
                    .collect(Collectors.toList());
        }

        // 如果还是没找到，使用第一个节点作为根
        if (roots.isEmpty()) {
            roots.add(list.get(0));
        }

        return buildTree(list, roots.get(0).getParentId());
    }

    /**
     * 扁平化树形结构
     * @param tree 树形结构
     * @return 扁平化后的列表
     */
    public static <T, N extends TreeNode<T, N>> List<N> flattenTree(List<N> tree) {
        List<N> result = new ArrayList<>();
        if (tree == null || tree.isEmpty()) {
            return result;
        }

        for (N node : tree) {
            result.add(node);
            if (node.getChildren() != null && !node.getChildren().isEmpty()) {
                result.addAll(flattenTree((List<N>) node.getChildren()));
            }
        }

        return result;
    }
}